﻿#include <ntddk.h>
#include <ntddser.h>

#define FILE_DEVICE_CUSTOM 0x8888
#define IOCTL_CUSTOM_HELLO CTL_CODE(FILE_DEVICE_CUSTOM, 0x800, METHOD_BUFFERED, FILE_ANY_ACCESS)


NTSTATUS DriverEntry(PDRIVER_OBJECT DriverObject, PUNICODE_STRING RegistryPath);
void     DriverUnload(PDRIVER_OBJECT DriverObject);
NTSTATUS VirtualSerialIoControl(PDEVICE_OBJECT DeviceObject, PIRP Irp);
NTSTATUS VirtualSerialCreate(PDEVICE_OBJECT DeviceObject, PIRP Irp);
NTSTATUS VirtualSerialClose(PDEVICE_OBJECT DeviceObject, PIRP Irp);

#define IOCTL_HELLO_DRIVER

NTSTATUS DriverEntry(PDRIVER_OBJECT DriverObject, PUNICODE_STRING RegistryPath)
{
    // DbgBreakPoint();

    UNICODE_STRING devName, symbolicLinkName;
    RtlInitUnicodeString(&devName, L"\\Device\\VirtualSerialPort");
    RtlInitUnicodeString(&symbolicLinkName, L"\\??\\VSP");
    PDEVICE_OBJECT pDeviceObject = NULL;

    DbgPrint("DriverEntry started.\n");

    NTSTATUS status = IoCreateDevice(DriverObject, 0, &devName, FILE_DEVICE_UNKNOWN, 0, FALSE, &pDeviceObject);
    if (!NT_SUCCESS(status))
    {
        DbgPrint("Failed to create device (status: %x)\n", status);
        return status;
    }

    // 设置DO_BUFFERED_IO标志
    pDeviceObject->Flags |= DO_BUFFERED_IO;

    status = IoCreateSymbolicLink(&symbolicLinkName, &devName);
    if (!NT_SUCCESS(status))
    {
        DbgPrint("Failed to create symbolic link (status: %x)\n", status);
        IoDeleteDevice(pDeviceObject);
        return status;
    }

    DriverObject->DriverUnload                         = DriverUnload;
    DriverObject->MajorFunction[IRP_MJ_CREATE]         = VirtualSerialCreate;
    DriverObject->MajorFunction[IRP_MJ_CLOSE]          = VirtualSerialClose;
    DriverObject->MajorFunction[IRP_MJ_DEVICE_CONTROL] = VirtualSerialIoControl;

    DbgPrint("DriverEntry completed successfully.\n");
    return STATUS_SUCCESS;
}

void DriverUnload(PDRIVER_OBJECT DriverObject)
{
    // DbgBreakPoint();

    DbgPrint("DriverUnload started.\n");

    UNICODE_STRING symbolicLinkName;
    RtlInitUnicodeString(&symbolicLinkName, L"\\??\\VSP");
    IoDeleteSymbolicLink(&symbolicLinkName);
    DbgPrint("Symbolic link deleted.\n");

    PDEVICE_OBJECT pDeviceObject = DriverObject->DeviceObject;
    while (pDeviceObject != NULL)
    {
        PDEVICE_OBJECT nextDevice = pDeviceObject->NextDevice;
        IoDeleteDevice(pDeviceObject);
        pDeviceObject = nextDevice;
    }

    DbgPrint("DriverUnload completed.\n");
}

NTSTATUS VirtualSerialCreate(PDEVICE_OBJECT DeviceObject, PIRP Irp)
{
    // DbgBreakPoint();

    Irp->IoStatus.Status      = STATUS_SUCCESS;
    Irp->IoStatus.Information = 0;
    IoCompleteRequest(Irp, IO_NO_INCREMENT);

    return STATUS_SUCCESS;
}

NTSTATUS VirtualSerialClose(PDEVICE_OBJECT DeviceObject, PIRP Irp)
{
    // DbgBreakPoint();
    Irp->IoStatus.Status      = STATUS_SUCCESS;
    Irp->IoStatus.Information = 0;
    IoCompleteRequest(Irp, IO_NO_INCREMENT);

    return STATUS_SUCCESS;
}

NTSTATUS VirtualSerialIoControl(PDEVICE_OBJECT DeviceObject, PIRP pIrp)
{
    // DbgBreakPoint();

    PVOID              pSysBuf = NULL;
    PIO_STACK_LOCATION pStack  = IoGetCurrentIrpStackLocation(pIrp);
    ULONG              ulCode  = pStack->Parameters.DeviceIoControl.IoControlCode;

    switch (ulCode)
    {
    case IOCTL_CUSTOM_HELLO:
	{
		pSysBuf = pIrp->AssociatedIrp.SystemBuffer;
		pIrp->IoStatus.Status = STATUS_SUCCESS;
		pIrp->IoStatus.Information = sizeof(ULONG);

		typedef struct _HelloInfor
		{
			ULONG uLineNumber;
			CHAR szHello[512];
		} HelloInfor, *PHelloInfor;

		PHelloInfor pHelloInfor = (PHelloInfor)ExAllocatePoolWithTag(NonPagedPool, sizeof(HelloInfor), 'info');
		if (pHelloInfor)
		{
			RtlCopyMemory(pHelloInfor, pSysBuf, sizeof(HelloInfor));
			DbgPrint("Kernel: %d %s", pHelloInfor->uLineNumber, pHelloInfor->szHello);
			ExFreePool(pHelloInfor);
		}
	}
        break;

    default:
        pIrp->IoStatus.Status = STATUS_INVALID_DEVICE_REQUEST;
        break;
    }

    IoCompleteRequest(pIrp, IO_NO_INCREMENT);
    return pIrp->IoStatus.Status;
}